from rag.rag_embeddings import OpenAIEmbeddings, LocalEmbeddings
from rag.rag_utils import FileLoader
from rag.rag_database import RAGDatabase
from tqdm import tqdm
import logging


# https://github.com/datawhalechina/happy-llm/blob/main/docs/chapter7/%E7%AC%AC%E4%B8%83%E7%AB%A0%20%E5%A4%A7%E6%A8%A1%E5%9E%8B%E5%BA%94%E7%94%A8.md
if __name__ == "__main__":
    file_loader = FileLoader("C:\JaredLyu\Project\MiniMind\\raw_data")
    files = file_loader.get_files()
    chunks = []
    for file in files:
        content = file_loader.get_file_content(file)
        file_chunks = file_loader.get_chunk(content)
        chunks.extend(file_chunks)
    
    openai_embeddings = OpenAIEmbeddings(model_name='BAAI/bge-large-zh', operator='huggingface')
    # local_embeddings = LocalEmbeddings(model_name_or_path="C:\JaredLyu\Project\MiniMind\MiniMind2")
    # local_embeddings = LocalEmbeddings(model_name_or_path="BAAI/bge-large-zh")
    rag_database = RAGDatabase(embeddings=openai_embeddings,use_cache=True)

    # logging.info(f"Starting to add {len(chunks)} chunks to the RAG database...")
    # for i, chunk in tqdm(enumerate(chunks)):
    #     rag_database.add_document(chunk)
    # rag_database.persist("C:\JaredLyu\Project\MiniMind\database\\rag_database.json")

    query = "东汉末年政治腐败吗？"
    results = rag_database.search(query, top_k=5)
    print(f"Query: {query}")
    for text, similarity in results:
        print(f"Text: {text}, Similarity: {similarity:.4f}")
    
